import torch
import torch.nn as nn
from torch.utils.data import DataLoader,Dataset
import albumentations as A
import numpy as np
import pandas as pd
import cv2
import matplotlib.pyplot as plt
from albumentations.pytorch import ToTensorV2
from Config import Config
from Model import Mymodel
train_transformer = A.Compose([
    A.Resize(height=Config.img_size, width=Config.img_size),
    A.OneOf([
        A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=0.5),
        #A.RandomBrightness(limit=0.1, p=0.5),
    ], p=1),
    # A.GaussNoise(),
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.5),
    A.RandomRotate90(p=0.5),
    A.ShiftScaleRotate(rotate_limit=1, p=0.5),
    # FancyPCA(alpha=0.1, p=0.5),
    # blur
    # A.OneOf([
    #     A.MotionBlur(blur_limit=3), A.MedianBlur(blur_limit=3), A.GaussianBlur(blur_limit=3),
    # ], p=0.5),
    # Pixels
    # A.OneOf([
    #     A.IAAEmboss(p=0.5),
    #     A.IAASharpen(p=0.5),
    # ], p=1),
    # Affine
    A.OneOf([
        A.ElasticTransform(p=0.5),
        #A.IAAPiecewiseAffine(p=0.5),
    ], p=1),
    #A.Normalize(mean=(0.5, 0.5,0.5),  max_pixel_value=255.0, p=1.0),
    ToTensorV2(p=1.0),
])
class Mydataset(Dataset):
    def __init__(self,path,transformer,mode = 'train'):
        self.path = path
        self.transformer = transformer
        self.mode = mode
        train_data = pd.read_csv(self.path+"train_data.csv",sep=',',names=["label","image"])
        self.len = len(train_data['image'][1:])
        self.image_path = self.path+"images/"+train_data['image'][1:]
        self.label = train_data['label'][1:]
    def __len__(self):
        return self.len


    def __getitem__(self, index):
        index = index+1
        image = self.image_path[index]
        label = self.label[index]
        if self.mode == "train":
            image = cv2.imread(image,cv2.COLOR_BGR2RGB)
            augmented = self.transformer(image=image,label=label)

            return augmented['image'].float(),torch.tensor(int(augmented['label'])).long()

class valdataset(Dataset):
    def __init__(self,val_path):
        self.path = val_path





if __name__ == '__main__':
    path = "/home/one/lhbdata/pet/pet_biometric_challenge_2022/train/"
    dataset = Mydataset(path,transformer=train_transformer,mode="train")
    dataloader = DataLoader(dataset,batch_size=1)
    model = Mymodel()
    valdataset = Mydataset(Config.val_path,transformer=train_transformer)
    for i,(img,label) in enumerate(dataloader):
        #print(label.data)
        output = model(img)
        _, preds = torch.max(output.data, 1)
        print(output.shape)
        #loss = TripletLoss(img, label)
        print(preds)
    # train_label= pd.read_csv(path+"train_data.csv",sep=',',names=['label','img'])
    # print(path+train_label['img'][1:])
    # print(train_label['label'][4])
